{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Convert YOLO labels to Label Studio format"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b200405be2acc893"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from ultralytics import YOLO\n",
    "\n",
    "MODEL_NAME = \"yolov8n.pt\"\n",
    "model = YOLO()\n",
    "\n",
    "def predict_yolo(images):\n",
    "    results = model(images)\n",
    "    predictions = []\n",
    "    for result in results:\n",
    "        img_width, img_height = result.orig_shape\n",
    "        boxes = result.boxes.cpu().numpy()\n",
    "        prediction = {'result': [], 'score': 0.0, 'model_version': MODEL_NAME}\n",
    "        scores = []\n",
    "        for box, class_id, score in zip(boxes.xywh, boxes.cls, boxes.conf):\n",
    "            x, y, w, h = box\n",
    "            prediction['result'].append({\n",
    "                'from_name': 'label',\n",
    "                'to_name': 'img',\n",
    "                'original_width': int(img_width),\n",
    "                'original_height': int(img_height),\n",
    "                'image_rotation': 0,\n",
    "                'value': {\n",
    "                    'rotation': 0,\n",
    "                    'rectanglelabels': [result.names[class_id]],\n",
    "                    'width': w / img_width * 100,\n",
    "                    'height': h / img_height * 100,\n",
    "                    'x': (x - 0.5 * w) / img_width * 100,\n",
    "                    'y': (y - 0.5 * h) / img_height * 100\n",
    "                },\n",
    "                'score': float(score),\n",
    "                'type': 'rectanglelabels',\n",
    "            })\n",
    "            scores.append(float(score))\n",
    "        prediction['score'] = min(scores) if scores else 0.0\n",
    "        predictions.append(prediction)\n",
    "    return predictions"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6ba1d477f4a9355d",
   "execution_count": 0
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Define Label Studio label config with YOLO labels"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e478ad6cf16a36e5"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "yolo_labels = '\\n'.join([f'<Label value=\"{label}\"/>' for label in model.names.values()])\n",
    "label_config = f'''\n",
    "<View>\n",
    "    <Image name=\"img\" value=\"$image\" zoom=\"true\" width=\"100%\" maxWidth=\"800\" brightnessControl=\"true\" contrastControl=\"true\" gammaControl=\"true\" />\n",
    "    <RectangleLabels name=\"label\" toName=\"img\">\n",
    "    {yolo_labels}\n",
    "    </RectangleLabels>\n",
    "</View>'''\n",
    "print(label_config)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6431c5ad6c16d75e",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Create Label Studio project"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "87121f74277b5104"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from label_studio_sdk.client import LabelStudio\n",
    "\n",
    "API_KEY = 'YOUR_API_KEY'\n",
    "client = LabelStudio(api_key=API_KEY)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fcd18afe7ce956de",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "project = client.projects.create(\n",
    "    title='Object detection',\n",
    "    description='Detect objects with YOLOv8',\n",
    "    label_config=label_config\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "dc584df69f390036",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can list all projects that are currently available in Label Studio, and check the project id."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ea4eb38e8ba4458b"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "for project in client.projects.list():\n",
    "    print(project.id, project.title)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1fabde7b2adac7da",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "In the future, we can get the project by its id."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "dc44bd2efd008d60"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "project = client.projects.get(id=28)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "966fc108d0c2dbc2",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "and update the project with new tasks"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ac2ceef6ab8a0103"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "task = client.tasks.create(\n",
    "    project=project.id,\n",
    "    data={\n",
    "    'image': 'https://hs-sandbox-pub.s3.amazonaws.com/images/SAM/bananas-on-a-table-picjumbo-com.jpg'\n",
    "    }\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "18296ffb62cd0248",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Import images from S3 bucket"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8293c41e5c1e3858"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "storage = client.import_storage.s3.create(\n",
    "    project=project.id,\n",
    "    bucket='hs-sandbox-pub',\n",
    "    prefix='images/shoppingmall',\n",
    "    regex_filter='.*jpg',\n",
    "    recursive_scan=True,\n",
    "    use_blob_urls=True,\n",
    "    aws_access_key_id='AKIAJZ5Q4ZQ7ZQ7ZQ7ZQ',\n",
    "    aws_secret_access_key='YOUR_SECRET_ACCESS_KEY',\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a1b589f7ebadc9d2",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "client.import_storage.s3.sync(id=storage.id)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "220a44b31f655fb3",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Listing all connected storages:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b246b58ecefb4d58"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "storage = client.import_storage.s3.list(project=project.id)[0]"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "927f7427efa37334",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Updating storage with new prefix and regex filter:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7b8841cc5c3da133"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "client.import_storage.s3.update(\n",
    "    id=storage.id,\n",
    "    project=project.id,\n",
    "    prefix='images/other',\n",
    "    regex_filter='.*jpg',\n",
    "    recursive_scan=True,\n",
    "    use_blob_urls=True\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "de8b0bf38d4afefe",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Import YOLO predictions to Label Studio"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5d867ee091c315f8"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "from tqdm import tqdm\n",
    "\n",
    "project = client.projects.get(28)\n",
    "tasks = client.tasks.list(project=project.id)\n",
    "images = []\n",
    "for i, task in enumerate(tqdm(tasks)):\n",
    "    url = f'http://localhost:8080{task[\"data\"][\"image\"]}'\n",
    "    image = Image.open(requests.get(url, headers={'Authorization': f'Token {API_KEY}'}, stream=True).raw)\n",
    "    predictions = predict_yolo([image])[0]\n",
    "    client.predictions.create(task=task['id'], result=predictions['result'], score=predictions['score'], model_version=predictions['model_version'])\n",
    "    if i > 10:\n",
    "        break"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "93e7d892aafa1877",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Managing batch of annotations\n",
    "\n",
    "We can use `views` to filter and organize tasks in Label Studio. For example, we can create a view that shows only tasks with low confidence predictions for the `person` class"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c0c6cc35c18f84c7"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "project = client.projects.get(28)\n",
    "for view in client.views.list(project=project.id):\n",
    "    print(view.id, view.data['title'])"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "946441699cbaa7fa",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "tab = client.views.create(\n",
    "    project=project.id,\n",
    "    data={\n",
    "        'title': 'Person low conf',\n",
    "        'filters': {\n",
    "            \"conjunction\": \"and\",\n",
    "            \"items\": [\n",
    "              {\n",
    "                \"filter\": \"filter:tasks:total_predictions\",\n",
    "                \"operator\": \"greater\",\n",
    "                \"value\": 0,\n",
    "                \"type\": \"Number\"\n",
    "              },\n",
    "              {\n",
    "                \"filter\": \"filter:tasks:predictions_results\",\n",
    "                \"operator\": \"contains\",\n",
    "                \"value\": \"person\",\n",
    "                \"type\": \"String\"\n",
    "              },\n",
    "              {\n",
    "                \"filter\": \"filter:tasks:predictions_score\",\n",
    "                \"operator\": \"less\",\n",
    "                \"value\": 0.3,\n",
    "                \"type\": \"Number\"\n",
    "              },\n",
    "            ]\n",
    "        }\n",
    "    }\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3e71a9dd1a235499",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now user can annotate only tasks that are filtered by the view, and mark them as `COMPLETED` when done."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "784f99b509ead277"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "client.views.update(\n",
    "    id=tab.id,\n",
    "    data={'title': 'COMPLETED'}\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "95b7504bf7ae3a12",
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Export annotations from Label Studio\n",
    "\n",
    "Finally, in order to export annotations that correspond to the batch of tasks, we can use the following code:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "44b0e1cb949e26d5"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "tab = client.views.get(id=36)\n",
    "annotated_tasks = client.tasks.list(view=tab.id, fields='all')\n",
    "for annotated_task in annotated_tasks:\n",
    "    print(annotated_task.annotations)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ccf01aff6ed5993b",
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
